package com.example.demo;

import jakarta.servlet.http.HttpServletRequest;
import org.apache.commons.io.FilenameUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Controller;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.SecureRandom;
import java.time.LocalDateTime;
import java.util.*;

@Controller
@RequestMapping("/upload")
public class FileUploadController {
    private static final Logger logger = LoggerFactory.getLogger(FileUploadController.class);
    // Whitelist of allowed extensions and MIME types
    private static final Map<String, String> ALLOWED_TYPES = Map.of(
            "jpg", "image/jpeg",
            "jpeg", "image/jpeg",
            "png", "image/png",
            "gif", "image/gif",
            "pdf", "application/pdf"
    );
    private static final long MAX_FILE_SIZE = 5 * 1024 * 1024; // 5MB
    private static final int MAX_FILES_PER_REQUEST = 5;
    private static final SecureRandom random = new SecureRandom();

    // Store uploads outside web root
    private static final String UPLOAD_ROOT = "uploads";

    @PostMapping(consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public ResponseEntity<String> handleFileUpload(@RequestParam("file") MultipartFile[] files,
                                                   HttpServletRequest request) {
        String user = request.getRemoteUser() != null ? request.getRemoteUser() : "anonymous";
        String ip = request.getRemoteAddr();
        if (files.length > MAX_FILES_PER_REQUEST) {
            logFailure("Too many files", user, ip, "file count", String.valueOf(files.length));
            return ResponseEntity.badRequest().body("Too many files");
        }
        List<String> savedFiles = new ArrayList<>();
        for (MultipartFile file : files) {
            String originalFilename = StringUtils.cleanPath(Objects.requireNonNull(file.getOriginalFilename()));
            // Validate file name
            if (!originalFilename.matches("^[a-zA-Z0-9._-]+$")) {
                logFailure("Invalid filename", user, ip, "filename", originalFilename);
                return ResponseEntity.badRequest().body("Invalid filename");
            }
            // Check file size
            if (file.getSize() > MAX_FILE_SIZE) {
                logFailure("File too large", user, ip, "size", String.valueOf(file.getSize()));
                return ResponseEntity.badRequest().body("File too large");
            }
            // Extension and MIME type check
            String ext = FilenameUtils.getExtension(originalFilename).toLowerCase();
            String allowedMime = ALLOWED_TYPES.get(ext);
            if (allowedMime == null) {
                logFailure("Disallowed extension", user, ip, "extension", ext);
                return ResponseEntity.badRequest().body("Disallowed file type");
            }
            String mimeType;
            try {
                mimeType = file.getContentType();
            } catch (Exception e) {
                logFailure("MIME detection failed", user, ip, "filename", originalFilename);
                return ResponseEntity.badRequest().body("MIME detection failed");
            }
            if (!allowedMime.equalsIgnoreCase(mimeType)) {
                logFailure("MIME mismatch", user, ip, "mime", mimeType);
                return ResponseEntity.badRequest().body("MIME type mismatch");
            }
            // File header (magic number) check
            try {
                if (!checkMagicNumber(file, ext)) {
                    logFailure("Magic number mismatch", user, ip, "filename", originalFilename);
                    return ResponseEntity.badRequest().body("File content does not match extension");
                }
            } catch (IOException e) {
                logFailure("Magic number check error", user, ip, "filename", originalFilename);
                return ResponseEntity.badRequest().body("File check error");
            }
            // Randomize file name and create subdirectory by hash
            String randomName = generateRandomName(ext);
            String subDir = randomName.substring(0, 2);
            Path uploadDir = Paths.get(UPLOAD_ROOT, subDir);
            try {
                Files.createDirectories(uploadDir);
            } catch (IOException e) {
                logFailure("Directory creation failed", user, ip, "dir", uploadDir.toString());
                return ResponseEntity.internalServerError().body("Server error");
            }
            Path filePath = uploadDir.resolve(randomName);
            try {
                // For images, re-encode to prevent webshells
                if (ext.equals("jpg") || ext.equals("jpeg") || ext.equals("png") || ext.equals("gif")) {
                    BufferedImage img = ImageIO.read(file.getInputStream());
                    if (img == null) {
                        logFailure("Image decode failed", user, ip, "filename", originalFilename);
                        return ResponseEntity.badRequest().body("Invalid image file");
                    }
                    ImageIO.write(img, ext.equals("jpg") ? "jpeg" : ext, filePath.toFile());
                } else {
                    file.transferTo(filePath);
                }
                savedFiles.add(filePath.toString());
            } catch (IOException e) {
                logFailure("File save error", user, ip, "filename", originalFilename);
                return ResponseEntity.internalServerError().body("File save error");
            }
        }
        return ResponseEntity.ok("Files uploaded: " + String.join(", ", savedFiles));
    }

    private void logFailure(String reason, String user, String ip, String param, String value) {
        logger.warn("[{}] [{}] [{}] {}: {}", LocalDateTime.now(), user, ip, param, value);
    }

    private String generateRandomName(String ext) {
        byte[] bytes = new byte[16];
        random.nextBytes(bytes);
        return Base64.getUrlEncoder().withoutPadding().encodeToString(bytes) + "." + ext;
    }

    private boolean checkMagicNumber(MultipartFile file, String ext) throws IOException {
        // Simple magic number check for demonstration
        byte[] header = new byte[8];
        file.getInputStream().read(header);
        if (ext.equals("jpg") || ext.equals("jpeg")) {
            return header[0] == (byte) 0xFF && header[1] == (byte) 0xD8;
        } else if (ext.equals("png")) {
            return header[0] == (byte) 0x89 && header[1] == (byte) 0x50 && header[2] == (byte) 0x4E && header[3] == (byte) 0x47;
        } else if (ext.equals("gif")) {
            return header[0] == (byte) 0x47 && header[1] == (byte) 0x49 && header[2] == (byte) 0x46;
        } else if (ext.equals("pdf")) {
            return header[0] == (byte) 0x25 && header[1] == (byte) 0x50 && header[2] == (byte) 0x44 && header[3] == (byte) 0x46;
        }
        return false;
    }
} 